from langchain.chains import RetrievalQA #检索QA链，在文档上进行检索
from langchain_community.document_loaders import CSVLoader #文档加载器，采用csv格式存储
from langchain_community.vectorstores import DocArrayInMemorySearch #向量存储
import pandas as pd
#导入向量存储索引创建器
from langchain.indexes import VectorstoreIndexCreator
from langchain_openai.embeddings import AzureOpenAIEmbeddings  # 导入嵌入模型
from openai import RateLimitError

from tool import get_azure_endpoint,get_api_key,get_api_version
import tenacity
if __name__ == '__main__':
    file = './OutdoorClothingCatalog_1000.csv'  # 使用langchain文档加载器对数据进行导入
    loader = CSVLoader(file_path=file)
    # 使用pandas导入数据，用以查看
    data = pd.read_csv(file, usecols=[1, 2])
    print(data.head())

    # 初始化嵌入模型（这里使用OpenAI的嵌入模型，你也可以选择其他模型）
    embeddings = AzureOpenAIEmbeddings(
        azure_endpoint=get_azure_endpoint().rstrip('/'),  # 移除尾部斜杠，只保留基础URL
        model="text-embedding-3-small",  # 重命名为 azure_deployment
        api_key=get_api_key(),
        api_version=get_api_version()
        )

    query = "请用markdown表格的方式列出所有具有防晒功能的衬衫，对每件衬衫描述进行总结"  # 使用索引查询创建一个响应，并传入这个查询


    @tenacity.retry(
        wait=tenacity.wait_fixed(60),  # Wait for 60 seconds before retrying
        retry=tenacity.retry_if_exception_type(RateLimitError),
        stop=tenacity.stop_after_attempt(3)  # Retry up to 3 times
    )
    def get_embedding(query):
        return embeddings.embed_query(query)
    vector = get_embedding(query)
    print(vector)
    # 创建指定向量存储类, 创建完成后，从加载器中调用, 通过文档加载器列表加载，一次性将这么多文档加载到服务器上容易超时
    index = VectorstoreIndexCreator(vectorstore_cls=DocArrayInMemorySearch,embedding=embeddings).from_loaders([loader])


    response = index.query(query)
    print(response)